from flask import Flask, request, jsonify
from flask_cors import CORS
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.prebuilt import create_react_agent
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
import os

# 初始化配置
os.environ["TAVILY_API_KEY"] = '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'
os.environ["DASHSCOPE_API_KEY"] = 'sk-5bc0688c5761427cadb9df012e589136'

# 创建大模型和数据库连接（全局对象）
llm = ChatTongyi(model="qwen-max")
db = SQLDatabase.from_uri('mysql+pymysql://root:123456@127.0.0.1:3306/mybatis_plus')
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()

# 系统提示词
system_prompt = """
您是一个被设计用来与洁客后台系统与数据库交互的智能助手，你的名字叫洁客超级智能体，设计你的人是研发组
的工程师令狐。
您的任务是：
1. **分析用户问题**，并创建一个 **正确的 SQL 查询**。
2. **执行 SQL 查询** 并返回相应的 **查询结果**。
3. **解读查询结果**，用 **自然语言** 向用户反馈答案。
4. **限制查询数据量**，除非用户明确指定，否则默认最多返回 10 条数据。
5. 允许你对数据库执行操作（INSERT、UPDATE、DELETE 等），操作成功以后要提醒我你操作成功了。
6. **先查看数据库中的表结构**，然后再执行查询，以确保 SQL 语句的正确性。
7.允许你
"""
# **不得执行任何修改数据库的操作（INSERT、UPDATE、DELETE 等）**。
# 创建 agent（全局对象）
prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content=system_prompt),
    MessagesPlaceholder(variable_name="messages")
])
agent_executor = create_react_agent(llm, tools, prompt=prompt)

# 创建 Flask 应用
app = Flask(__name__)
CORS(app)  # 允许跨域请求


@app.route('/query', methods=['POST'])
def handle_query():
    # 获取请求参数
    data = request.json
    if not data or 'content' not in data:
        return jsonify({"error": "Missing 'content' parameter"}), 400

    try:
        # 执行查询
        resp = agent_executor.invoke({
            "messages": [HumanMessage(content=data['content'])]
        })

        # 处理返回结果
        result = resp['messages']
        if len(result) < 1:
            return jsonify({"error": "No response generated"}), 500

        return jsonify({
            "answer": result[-1].content,
            "full_messages": [msg.dict() for msg in result]
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)